#include "quick_launch_win.h"

#include <atlenc.h>
#include <ShlObj.h>

#include "ximage/ximage.h"

#include "icons.h"
#include "utils.h"

namespace quick_launch {

WinVersion GetWinVersion() {
  static bool checked_version = false;
  static WinVersion win_version = WINVERSION_PRE_2000;
  if (!checked_version) {
    OSVERSIONINFOEX version_info;
    version_info.dwOSVersionInfoSize = sizeof version_info;
    GetVersionEx(reinterpret_cast<OSVERSIONINFO*>(&version_info));
    if (version_info.dwMajorVersion == 5) {
      switch (version_info.dwMinorVersion) {
    case 0:
      win_version = WINVERSION_2000;
      break;
    case 1:
      win_version = WINVERSION_XP;
      break;
    case 2:
    default:
      win_version = WINVERSION_SERVER_2003;
      break;
      }
    } else if (version_info.dwMajorVersion == 6) {
      if (version_info.wProductType != VER_NT_WORKSTATION) {
        // 2008 is 6.0, and 2008 R2 is 6.1.
        win_version = WINVERSION_2008;
      } else {
        if (version_info.dwMinorVersion == 0) {
          win_version = WINVERSION_VISTA;
        } else {
          win_version = WINVERSION_WIN7;
        }
      }
    } else if (version_info.dwMajorVersion > 6) {
      win_version = WINVERSION_WIN7;
    }
    checked_version = true;
  }
  return win_version;
}

struct Browser_Param {
  TCHAR initial_path[MAX_PATH];
  TCHAR title[MAX_PATH];
};


// The callback function for select local path and
// initialize the inital path.
int WINAPI BrowserCallBack(HWND hwnd, UINT uMsg, LPARAM lParam, 
                           LPARAM lpData) {
  switch (uMsg) {
    case BFFM_INITIALIZED:
      Browser_Param* param = (Browser_Param*)lpData;
      SendMessage(hwnd, BFFM_SETSELECTION, TRUE, (LPARAM)param->initial_path);
      SetWindowText(hwnd, param->title);
      break;
  }

  return 0;
}

bool SaveLink(HWND hwnd, LPCTSTR filepath, LPCTSTR exefile, 
              LPCTSTR iconfile, int iconindex, LPCTSTR params) {
  HRESULT hres;
  IShellLink* psl;

  // Create IShellink interface.
  hres = CoCreateInstance(CLSID_ShellLink, NULL,
                          CLSCTX_INPROC_SERVER, IID_IShellLink, (void**)&psl); 

  if (SUCCEEDED(hres)) {
    IPersistFile* ppf;
    // Query IPersistFile interface.
    hres = psl->QueryInterface(IID_IPersistFile, (void**)&ppf);
    if (SUCCEEDED(hres)) {
      // Get work path corresponding to exe file name.
      TCHAR workpath[MAX_PATH];
      _tcscpy(workpath, exefile);
      TCHAR* index = _tcsrchr(workpath, '\\');
      if (index) {
        *index = 0;
      }
      // Set parameters.
      psl->SetPath(exefile);
      psl->SetArguments(params);
      psl->SetIconLocation(iconfile, iconindex);
      psl->SetWorkingDirectory(workpath);
      // Save link to disk.
      hres = ppf->Save(filepath, TRUE);
      ppf->Release();
    }
    psl->Release();
  }

  // Return the result, S_OK mean success.
  return hres == S_OK ? true : false; 
}

bool ResolveLink(HWND hwnd, LPCTSTR lpszLinkFile, std::string& path, 
                 std::string& params, std::string& iconfile, int& iconindex,
                 std::string& icondata, bool resolve_icon) {
  HRESULT hres;
  IShellLink* psl;
  TCHAR temp_path[MAX_PATH];
  TCHAR exe_file[MAX_PATH];
  TCHAR paramters[INFOTIPSIZE];

  WIN32_FIND_DATA wfd;

  // Create IShellLink interface.
  hres = CoCreateInstance(CLSID_ShellLink, NULL,
                          CLSCTX_INPROC_SERVER, IID_IShellLink, (void**)&psl);
  if (SUCCEEDED(hres)) {
    IPersistFile* ppf;

    // Query IPersistFile interface.
    hres = psl->QueryInterface(IID_IPersistFile, (void**)&ppf);
    if (SUCCEEDED(hres)) {
      // Load link file and resolve it.
      hres = ppf->Load(lpszLinkFile, STGM_READ);
      if(FAILED(hres))
        return false;
      hres = psl->Resolve(hwnd, SLR_NO_UI);
      if (FAILED(hres)) 
        return false;

      // Get the exe file.
      hres = psl->GetPath(temp_path, MAX_PATH, (WIN32_FIND_DATA*)&wfd,
                          SLGP_SHORTPATH);
      GetLongPathName(temp_path, exe_file, INFOTIPSIZE);
      path = utils::TCHARToUtf8(exe_file).get_utf8();

      // Resolve icon data. if the icon file is a ico file, 
      // get icon data directly, otherwise extract the icon data.
      if (resolve_icon) {
        TCHAR icon_file[MAX_PATH] = _T("");
        int icon_index = 0;
        hres = psl->GetIconLocation(temp_path, MAX_PATH, &icon_index);

        // Some time the icon location file is empty.
        // So we need copy the exe_file to the temp_path.
        if (_tcslen(temp_path) == 0)
          memcpy(temp_path, exe_file, sizeof(temp_path));
        
        // some time the icon location include a environment string such as
        // %USERPROFILE%, so we need resolve it to the real path.
        ExpandEnvironmentStrings(temp_path, icon_file, MAX_PATH);
        iconindex = icon_index;

        if (_taccess(icon_file, 0) != 0) 
          _tcscpy(icon_file, _T(""));

        if (SUCCEEDED(hres)) {
          GetIconData(icon_file, icon_index, icondata);
        }
        iconfile = utils::TCHARToUtf8(icon_file).get_utf8();
      }

      hres = psl->GetArguments(paramters, INFOTIPSIZE);
      params = utils::TCHARToUtf8(paramters).get_utf8();
      ppf->Release();   
    }
    psl->Release();
  }

  // Return the result, S_OK mean success.
  return hres == S_OK ? true : false;
}

bool DeleteQuickLaunch(HWND hwnd, const char* path) {
  utils::Utf8ToTCHAR file_name(path);
  std::string exepath, params, iconfile, icondata;
  int iconindex;

  // Resolve the link, we need delete the user data directory.
  ResolveLink(hwnd, file_name, exepath, params, iconfile, 
              iconindex, icondata, false);

  // Get the user data dir.
  TCHAR user_data_dir[MAX_PATH] = _T("");
  ResolveUserDataDir(params, user_data_dir);
  utils::Utf8ToTCHAR exe_file(exepath.c_str());
  char* default_path = GetDefaultDataPath(exe_file);
  if (_tcscmp(user_data_dir, utils::Utf8ToTCHAR(default_path)) == 0) {
    NPN_MemFree(default_path);
    if (DeleteFile(file_name))
      return true;
    else
      return false;
  }

  if (_taccess(file_name, 0) != 0)
    return true;

  if (DeleteFile(file_name))
    return true;
  else
    return false;
}

bool DeleteQuickLaunch(HWND hwnd, TCHAR* path) {
  utils::TCHARToUtf8 link_file(path);
  return DeleteQuickLaunch(hwnd, link_file.get_utf8());
}

bool DeleteQuickLaunch(TCHAR* path) {
  if (DeleteFile(path))
    return true;
  else
    return false;
}

char* GetLocalPath(HWND hwnd, const char* title, const char* inital_path) {
  TCHAR display_name[MAX_PATH] = {0};
  Browser_Param param = {0};

  if (inital_path != NULL) {
    utils::Utf8ToTCHAR temp(inital_path);
    _tcscpy(param.initial_path, temp);
  }
  if (title != NULL) {
    utils::Utf8ToTCHAR temp(title);
    _tcscpy(param.title, temp);
  }

  BROWSEINFO info={0};
  info.hwndOwner = hwnd;
  info.lpszTitle = param.title;
  info.pszDisplayName = display_name;
  info.lpfn = BrowserCallBack;
  info.ulFlags = BIF_RETURNONLYFSDIRS | BIF_NEWDIALOGSTYLE;
  info.lParam = (LPARAM)&param;
  BOOL ret = SHGetPathFromIDList(SHBrowseForFolder(&info), display_name);

  utils::TCHARToUtf8 folder(ret ? display_name : param.initial_path);
  int size = strlen(folder.get_utf8());
  char* return_path = (char*)NPN_MemAlloc(size + 1);
  strcpy(return_path, folder.get_utf8());
  return return_path;
}

bool SelectQuickLaunch(HWND hwnd, TCHAR* quick_launch, int size) {
  OPENFILENAME ofn = {0};
  ofn.lStructSize = sizeof(OPENFILENAME);
  ofn.hwndOwner = hwnd;
  ofn.lpstrFilter = _T("Quick Lanuch\0*.lnk\0");
  ofn.lpstrFile = quick_launch;
  ofn.nMaxFile = size;
  ofn.lpstrFileTitle = NULL;
  ofn.nMaxFileTitle = 0;
  ofn.lpstrInitialDir = NULL;
  ofn.Flags = OFN_SHOWHELP | OFN_OVERWRITEPROMPT;
  ofn.lpstrTitle = _T("Select Chrome Lanuch");
  ofn.lpstrDefExt = _T("lnk");
  if (GetOpenFileName(&ofn))
    return true;
  else
    return false;
}

bool Startup(HWND hwnd, const char* filename) {
  utils::Utf8ToTCHAR linkfile(filename);
  std::string path;
  std::string params;
  std::string icondata;
  std::string iconfile;
  int iconindex;

  if (ResolveLink(hwnd, linkfile, path, params, 
      iconfile, iconindex, icondata)) {
    int index = params.find("--load-extension=\"");
    int offset = strlen("--load-extension=\"");
    char end_char = '"';
    
    TCHAR theme_helper_dir[MAX_PATH] = _T("");
    TCHAR manifest_json[MAX_PATH] = _T("");
    std::string load_extension;  
    if (index != params.npos) {
      int start_index = index + offset;
      index = start_index;
      int len = params.length();
      while (index < len && params[index] != end_char)
        index++;
      load_extension = params.substr(start_index, index - start_index);
      _tcscpy(theme_helper_dir, utils::Utf8ToTCHAR(load_extension.c_str()));
      _stprintf(manifest_json, _T("%s\\%s"), theme_helper_dir, 
                _T("manifest.json"));
      if (_taccess(manifest_json, 0) != 0) {
        start_index -= offset;
        params.erase(start_index, index - start_index + 1);
        SaveLink(hwnd, linkfile, utils::Utf8ToTCHAR(path.c_str()), 
                 utils::Utf8ToTCHAR(iconfile.c_str()), iconindex, 
                 utils::Utf8ToTCHAR(params.c_str()));
      }
    }
  }

  HINSTANCE ret = ShellExecute(hwnd, _T("open"),
                               linkfile, NULL, NULL, SW_SHOWDEFAULT);
  if ((int)ret > 32)
    return true;
  else
    return false;
}

bool ValidPath(const char* path) {
  utils::Utf8ToTCHAR path_name(path);

  DWORD attr = GetFileAttributes(path_name);
  if (attr != INVALID_FILE_ATTRIBUTES && (attr & FILE_ATTRIBUTE_DIRECTORY))
    return true;
  else
    return false;
}

bool ValidPath(TCHAR* path) {
  DWORD attr = GetFileAttributes(path);
  if (attr != INVALID_FILE_ATTRIBUTES && (attr & FILE_ATTRIBUTE_DIRECTORY))
    return true;
  else
    return false;
}

char* GetDesktopFolder(HWND hwnd) {
  TCHAR desktop_folder[MAX_PATH] = _T("");
  SHGetSpecialFolderPath(hwnd, desktop_folder, CSIDL_DESKTOPDIRECTORY, 0);
  utils::TCHARToUtf8 folder(desktop_folder);
  int size = strlen(folder.get_utf8());
  char* return_value = (char*)NPN_MemAlloc(size + 1);
  strcpy(return_value, folder.get_utf8());
  return return_value;
}

char* GetQuickLaunchFolder(HWND hwnd) {
  TCHAR quick_launch_folder[MAX_PATH] = _T("");
  SHGetSpecialFolderPath(hwnd, quick_launch_folder, CSIDL_APPDATA, 0);
  _tcscat(quick_launch_folder, 
          _T("\\Microsoft\\Internet Explorer\\Quick Launch"));

  utils::TCHARToUtf8 folder(quick_launch_folder);
  int size = strlen(folder.get_utf8());
  char* return_value = (char*)NPN_MemAlloc(size + 1);
  strcpy(return_value, folder.get_utf8());
  return return_value;
}

char* GetDefaultDataPath(TCHAR* exefile) {
  TCHAR exe_file[MAX_PATH] = _T("");
  if (exefile == NULL)
    GetModuleFileName(GetModuleHandle(NULL), exe_file, MAX_PATH);
  else
    _tcscpy(exe_file, exefile);
  TCHAR* backblash = _tcsrchr(exe_file, '\\');
  int loopcount = 2;
  if (backblash) {
    do {
      *backblash = 0;
      backblash = _tcsrchr(exe_file, '\\');
    } while (backblash && --loopcount);
  }
  _tcscat(exe_file, _T("\\User Data"));
  utils::TCHARToUtf8 folder(exe_file);
  int size = strlen(folder.get_utf8());
  char* return_value = (char*)NPN_MemAlloc(size + 1);
  strcpy(return_value, folder.get_utf8());
  return return_value;
}

char* GetUserDataFolder() {
  TCHAR user_data_dir[MAX_PATH] = _T("");
  if (GetWinVersion() >= WINVERSION_VISTA) {
    SHGetSpecialFolderPath(NULL, user_data_dir, CSIDL_LOCAL_APPDATA, FALSE);
    _tcscat(user_data_dir, _T("\\ChromeUserData"));
  } else {
    TCHAR drivers[200] = _T("");
    TCHAR* max_free_space_dir = NULL;
    TCHAR* p = NULL;
    GetLogicalDriveStrings(200, drivers);
    p = drivers;
    ULARGE_INTEGER max_free_bytes = {0};
    do {
      ULARGE_INTEGER totalfreebytes = {0};
      GetDiskFreeSpaceEx(p, NULL, NULL, &totalfreebytes);
      if (totalfreebytes.QuadPart > max_free_bytes.QuadPart) {
        max_free_bytes = totalfreebytes;
        max_free_space_dir = p;
      }
      while(*p++);
    } while (*p);

    if (max_free_space_dir) {
      if (max_free_space_dir[_tcslen(max_free_space_dir) - 1] == '\\')
        max_free_space_dir[_tcslen(max_free_space_dir) - 1] = 0;
      _stprintf(user_data_dir, _T("%s\\ChromeUserData"), max_free_space_dir);
      if (!ValidPath(user_data_dir))
        CreateDirectory(user_data_dir, NULL);
    } else {
      return NULL;
    }
  }
  utils::TCHARToUtf8 folder(user_data_dir);
  int size = strlen(folder.get_utf8());
  char* return_value = (char*)NPN_MemAlloc(size + 1);
  strcpy(return_value, folder.get_utf8());
  return return_value;
}

bool CreateFolder(TCHAR* path, int size) {
  for (int index = 0; index < size; index++) {
    if (path[index] == '\\') {
      path[index] = 0;
      if (!ValidPath(path)) {
        CreateDirectory(path, NULL);
      }
      path[index] = '\\';
    }
  }

  if (!ValidPath(path))
    return CreateDirectory(path, NULL);
  else
    return true;
}

bool CopyFolder(TCHAR* source, TCHAR* dest) {
  TCHAR from[MAX_PATH] = _T("");
  TCHAR to[MAX_PATH] = _T("");
  _tcscpy(from, source);
  _tcscpy(to, dest);

  SHFILEOPSTRUCT file_op; 
  file_op.hwnd = NULL; 
  file_op.wFunc = FO_COPY; 
  file_op.pFrom = from; 
  file_op.pTo = to; 
  file_op.fFlags = FOF_NOCONFIRMATION | FOF_NOCONFIRMMKDIR | 
      FOF_SILENT | FOF_NOERRORUI;
  file_op.hNameMappings = NULL; 
  file_op.lpszProgressTitle = NULL;
  if (!SHFileOperation(&file_op)) 
    return true;
  else
    return false;
}

TCHAR* GetSystemIconsPath(HWND hwnd, TCHAR* path) {
  if (path) {
    SHGetSpecialFolderPath(hwnd, path, CSIDL_MYDOCUMENTS, FALSE);
    _tcscat(path, _T("\\ExtensionIcons\\SystemIcons\\"));
  }
  return path;
}

TCHAR* GetUserDefinedIconsPath(HWND hwnd, TCHAR* path) {
  if (path) {
    SHGetSpecialFolderPath(hwnd, path, CSIDL_MYDOCUMENTS, FALSE);
    _tcscat(path, _T("\\ExtensionIcons\\UserDefined\\"));
  }
  return path;
}

void GetIconData(TCHAR* iconfile, int index, std::string& icondata) {
  TCHAR extname[10];
  _tsplitpath(iconfile, NULL, NULL, NULL, extname);

  // Get the icon data from the ico file directly.
  if (_tcsicmp(extname, _T(".ico")) == 0) {
    FILE* f = _tfopen(iconfile, _T("rb"));
    if (f) {
      fseek(f, 0, SEEK_END);
      int len = ftell(f);
      fseek(f, 0, SEEK_SET);
      byte* source = (byte*)malloc(len);
      if (source) {
        int read_len = 0;
        while ((read_len += fread(
          source + read_len, 1, len - read_len, f)) != len) ;
        int dest_len = Base64EncodeGetRequiredLength(len);
        char* base64 = (char*)malloc(dest_len + 1);
        if (base64) {
          Base64Encode(source, len, base64, 
                       &dest_len, ATL_BASE64_FLAG_NOCRLF);
          base64[dest_len] = 0;
          icondata += base64;
          free(base64);
          free(source);
        }
      }
      fclose(f);
    }
  } else {
    // Extract the icon data from exe or dll file.
    IconFile icon_file = Icons::ExtractExeIcon(iconfile, index);
    if (icon_file.datasize > 0 && icon_file.data) {
      int dest_len = Base64EncodeGetRequiredLength(icon_file.datasize);
      char* base64 = (char*)malloc(dest_len + 1);
      if (base64) {
        Base64Encode(icon_file.data, icon_file.datasize, base64, 
                     &dest_len, ATL_BASE64_FLAG_NOCRLF);
        base64[dest_len] = 0;
        icondata += base64;
        free(base64);
      }
    }
    Icons::FreeData(icon_file);
  }
}

bool ConvertImageToIcon(TCHAR* imagefile, TCHAR* iconfile) {
  FILE* source = _tfopen(imagefile, _T("rb"));
  if (source) {
    CxImage image;
    if (image.Decode(source, CXIMAGE_FORMAT_UNKNOWN)) {
      FILE* dest = _tfopen(iconfile, _T("wb"));
      image.Resample(48, 48);
      if (dest && image.Encode(dest, CXIMAGE_FORMAT_ICO)) {
        fclose(source);
        fclose(dest);
        return true;
      }
    }
    fclose(source);
  }

  return false;
}

TCHAR* ResolveUserDataDir(const std::string& params, TCHAR* userdatadir) {
  int index = params.find("--user-data-dir=\"");
  int offset = strlen("--user-data-dir=\"");
  char end_char = '"';
  if (index == params.npos) {
    index = params.find("--user-data-dir=");
    end_char = ' ';
    offset--;
  }

  std::string user_data_dir;  
  if (index != params.npos) {
    int start_index = index + offset;
    index = start_index;
    int len = params.length();
    while (index < len && params[index] != end_char)
      index++;
    user_data_dir = params.substr(start_index, index - start_index);
  }
  _tcscpy(userdatadir, utils::Utf8ToTCHAR(user_data_dir.c_str()));
  return userdatadir;
}

WCHAR* ResolveUserDataDir(const std::wstring& params, WCHAR* userdatadir) {
  int index = params.find(L"--user-data-dir=\"");
  int offset = wcslen(L"--user-data-dir=\"");
  WCHAR end_char = '"';
  if (index == params.npos) {
    index = params.find(L"--user-data-dir=");
    end_char = ' ';
    offset--;
  }

  std::wstring user_data_dir;
  if (index != params.npos) {
    int start_index = index + offset;
    index = start_index;
    int len = params.length();
    while (index < len && params[index] != end_char)
      index++;
    user_data_dir = params.substr(start_index, index - start_index);
  }
  wcscpy(userdatadir, user_data_dir.c_str());
  return userdatadir;
}

}